package ch.epfl.lara.matcheck.verify.transform;

import ch.epfl.lara.matcheck.ast.{Formula,And,Tru,Imply,Neg}
import ch.epfl.lara.matcheck.util.ShapeKeeper
import scala.util.parsing.input.Positional

/**
 * Refine each program point by using context information.
 * This transformation allow finer checks of pattern matching.
 */
object VMatchContextRefinement extends ShapeKeeper {
  
  def apply[T<:VTree with Positional](vtree: T): T = (vtree match {
    case VProgram(constr,conts)                 => cast[T,VProgram]   (VProgram(constr,conts.map(this(_))))
    case VContainer(inv,methods)                     => cast[T,VContainer] (VContainer(inv,methods.map(this(_))))
    case VMethodDef(bindings,ppc @ Some((pre,_)),body) => cast[T,VMethodDef] (VMethodDef(bindings,ppc,this((Tru,pre),body)))
    case VMethodDef(bindings,ppc @ None,body)        => cast[T,VMethodDef] (VMethodDef(bindings,ppc,this(body)))
    case VIf(cond,then,elze)                         => cast[T,VIf]        (VIf(cond,
                                                                                this(cond,then),
                                                                                this((cond._1,Neg(cond._2)),elze))
                                                                           ) 
    case vm: VMatch                                  => cast[T,VMatch]     (refine((Tru,Tru),vm))
    case _                                           => vtree
  }).setPos(vtree.pos)

  def apply[T<:VExpr with Positional](refinement: (Formula,Formula), vexpr: T): T = (vexpr match {
    case VIf(cond,then,elze) => cast[T,VIf]  (VIf(cond, 
                                                  this( (And(refinement._1,cond._1),
                                                         And(refinement._2,cond._2)),then
                                                       ),
                                                  this( (And(refinement._1,cond._1),
                                                         And(refinement._2,Neg(cond._2))),elze
                                                       ))
                                              )
    case vm: VMatch          => cast[T,VMatch] (refine(refinement,vm))
    case _                   => vexpr 
  }).setPos(vexpr.pos)
  
  
  private def refine(refinement: (Formula,Formula), vm: VMatch): VMatch = {
    val VMatch(scrut,cases) = vm
    val (decl,form) = refinement
    var nested_refinement: (Formula,Formula) = (And(decl,scrut),form)
    val vcases = for(VCaseDef(pos,pFormula @ Imply(lhs,rhs),body)<-cases) yield {
      val current_refinement = (And(nested_refinement._1,lhs),And(nested_refinement._2,rhs))
      val vcase = VCaseDef(pos,pFormula,this(current_refinement,body))
      nested_refinement = (And(nested_refinement._1,lhs),And(nested_refinement._2,Neg(rhs)))
      vcase
    }
    VMatch(And(And(decl,form),scrut),vcases)
  }
}
